#pragma once
#include <iostream>
#include <utility>
#include <exception>
#include <stack>

template<typename Key, typename T>
struct TreeNode {
    std::pair<Key, T> data;
    TreeNode* left;
    TreeNode* right;
    TreeNode* parent;
    TreeNode(const Key& key, const T& firstue, TreeNode* parentNode = nullptr):
        data(std::make_pair(key, firstue)), left(nullptr), right(nullptr), parent(parentNode) {}
};

template<typename Key, typename T>
class MyMap {
public:
    MyMap():root_(nullptr) {}

    // 回收所有节点
    ~MyMap() {
        clear(this->root_);
    }

    MyMap(const MyMap& other) = delete;
    MyMap& operator=(const MyMap& other) = delete;

    void insert(const Key& key, const T& firstue) {
        if (nullptr == root_) {
            root_ = new TreeNode<Key, T>(key, firstue);
            return;
        }

        TreeNode<Key, T> * current = root_;
        TreeNode<Key, T> * parent = nullptr;
        // 找到需要插入的位置
        while (nullptr != current) {
            if (key < current->data.first) {
                parent = current;
                current = current->left;
            } else if (key > current->data.first) {
                parent = current;
                current = current->right;
            } else {
                current->data.second = firstue;
                return;
            }
        }
        if (key < parent->data.first) {
            parent->left = new TreeNode<Key, T>(key, firstue, parent);
        } else {
            parent->right = new TreeNode<Key, T>(key, firstue, parent);
        }
    }

    // 删除某个结点，需要找到对应的前驱或后继
    void erase(const Key& key) {
        // 找到该结点
        auto * node = find(key);
        if (nullptr == node) return;

        // 情况1:结点有两个子树（用后继结点替换)
        if (nullptr != node->left && nullptr != node->right) {
            TreeNode<Key, T>* succ = successor(node); // 找后继
            // 用后继节点覆盖当前节点
            node->data = succ->data;
            // 转为删除后继节点
            node = succ;
        }

        auto * child = (nullptr != node->left) ? node->left : node->right;
        if (nullptr != child) {
            child->parent = node->parent;
        }

        if (nullptr == node->parent) { // 删除的是根节点
            root_ = child;
        } else if (node == node->parent->left) {
            node->parent->left = child;
        } else {
            node->parent->right = child;
        }
        delete node;
    }   

    TreeNode<Key, T> * find(const Key& key) const {
        auto * current = this->root_;
        while(nullptr != current) {
            if (key < current->data.first) {
                current = current->left;
            } else if (key > current->data.first) {
                current = current->right;
            } else {
                return current;
            }
        } 
        return nullptr;
    }

    void clear() {
        clear(this->root_);
        root_ = nullptr;
    }

    class Iterator {
    public:
        Iterator(TreeNode<Key, T>* node = nullptr) : current(node){}
        std::pair<const Key, T>& operator*() const { return current->data; }
        std::pair<Key, T>* operator->() const { return &current->data; }
        bool operator==(const Iterator& other) const {
            return current == other.current;
        }

        bool operator!=(const Iterator& other) const {
            return current != other.current;
        }
        // 前置递增
        Iterator& operator++() {
            current = successor(current);
            return *this;
        }

        // 后置递增
        Iterator operator++(int) {
            Iterator temp = *this;
            current = successor(current);
            return temp;
        }
    private:
        TreeNode<Key, T> * current;
        TreeNode<Key, T> * minimum(TreeNode<Key, T> * node) const {
            if (node == nullptr) return nullptr;
            while (nullptr != node->left) 
            {
                node = node->left;
            }
            return node;
        }

        TreeNode<Key, T> * successor(TreeNode<Key, T> * node) const {
            // 有后继，先找后继结点
            // 右子树 最左下结点
            if (nullptr != node->right) {
                return minimum(node->right);
            }
            // 没有后继，先找前驱
            // 左子树 最右下结点
            auto * p = node->parent;
            while(nullptr != p && node == p->right) {
                node = p;
                p = p->parent;
            }
            return p;
        }
    };

    Iterator begin() const {
        return Iterator(minimum(root_));
    }

    Iterator end() const {
        return Iterator(nullptr);
    }

private:
    void clear(TreeNode<Key, T> *node) {
        if (nullptr == node) return;
        clear(node->left);
        clear(node->right);
        delete node;
    }

    // 
    TreeNode<Key, T> * minimum(TreeNode<Key, T> * node) const {
        if (node == nullptr) return nullptr;
        while (nullptr != node->left) 
        {
            node = node->left;
        }
        return node;
    }

    TreeNode<Key, T> * maximum(TreeNode<Key, T> * node) const {
        if (node == nullptr) return nullptr;
        while (nullptr != node->right) 
        {
            node = node->right;
        }
        return node;
    }

    TreeNode<Key, T> * successor(TreeNode<Key, T> * node) const {
        // 有后继，先找后继结点
        // 右子树 最左下结点
        if (nullptr != node->right) {
            return minimum(node->right);
        }
        // 没有后继，先找前驱
        // 左子树 最右下结点
        auto * p = node->parent;
        while(nullptr != p && node == p->right) {
            node = p;
            p = p->parent;
        }
        return p;
    }

    TreeNode<Key, T> * root_;
};